package cmd

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"fmt"
	"io/ioutil"

	"gitee.com/qianyucc/secure-tool/util"
	"github.com/spf13/cobra"
)

var (
	PublicKey       string
	PrivateKey      string
	GenerateKeySize int64
	RsaKeyFileName  string = "key"
)

// rsaCmd 16进制编码
var rsaCmd = &cobra.Command{
	Use:   "rsa",
	Short: "Encrypt and decrypt files and strings using RSA algorithm",
	Long:  `Encrypt and decrypt files and strings using RSA algorithm`,
	Run: func(cmd *cobra.Command, args []string) {
		if err := cmd.ParseFlags(args); err != nil {
			fmt.Println(err)
			return
		}
		var raw []byte
		var err error
		if IsDecode {
			if Text == "" {
				fmt.Println("invalid parameters, please enter the text to decrypt.")
				return
			}
			if PrivateKey == "" {
				fmt.Println("invalid parameters, please enter the private key file name.")
				return
			}
			var content []byte
			if content, err = base64.StdEncoding.DecodeString(Text); err != nil {
				fmt.Println(err)
				return
			}
			if raw, err = DecryptTextWithRsaByPrivateKey(content, PrivateKey); err != nil {
				fmt.Println(err)
				return
			}
			util.PrintAndCopyTextToClipboardIgnoreError(string(raw))
		} else {
			// 生成公钥和私钥
			if GenerateKeySize != -1 {
				if err := GenerateRsaKey(RsaKeyFileName, GenerateKeySize); err != nil {
					fmt.Println(err)
					return
				}
				return
			}
			// 加密
			if Text == "" {
				fmt.Println("invalid parameters, please enter the text to encrypt.")
				return
			}
			if PublicKey == "" {
				fmt.Println("invalid parameters, please enter the public key file name.")
				return
			}
			if raw, err = EncryptTextWithRsaByPublicKey([]byte(Text), PublicKey); err != nil {
				fmt.Println(err)
				return
			}
			util.PrintAndCopyTextToClipboardIgnoreError(base64.StdEncoding.EncodeToString(raw))
		}
	},
}

// GenerateRsaKey 生成RSA公钥和私钥
func GenerateRsaKey(fileName string, size int64) (err error) {
	privateKey, err := rsa.GenerateKey(rand.Reader, int(size))
	if err != nil {
		return err
	}
	publicKey := privateKey.PublicKey
	privateKeyPerm := pem.EncodeToMemory(
		&pem.Block{
			Type:  "RSA PRIVATE KEY",
			Bytes: x509.MarshalPKCS1PrivateKey(privateKey),
		},
	)
	derPkix, err := x509.MarshalPKIXPublicKey(&publicKey)
	if err != nil {
		return err
	}
	publicKeyPerm := pem.EncodeToMemory(
		&pem.Block{
			Type:  "RSA PUBLIC KEY",
			Bytes: derPkix,
		},
	)
	if err := ioutil.WriteFile("./"+fileName+".rsa", privateKeyPerm, 0700); err != nil {
		return err
	}
	if err := ioutil.WriteFile("./"+fileName+".rsa.pub", publicKeyPerm, 0700); err != nil {
		return err
	}
	return nil
}

// EncryptTextWithRsaByPublicKey 公钥加密字符串
func EncryptTextWithRsaByPublicKey(content []byte, publicKeyFileName string) (raw []byte, err error) {
	publicKeyRaw, err := ioutil.ReadFile(publicKeyFileName)
	if err != nil {
		return nil, err
	}
	block, _ := pem.Decode(publicKeyRaw)
	publicKey, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return nil, err
	}
	return rsa.EncryptPKCS1v15(rand.Reader, publicKey.(*rsa.PublicKey), content)
}

// DecryptTextWithRsaByPrivateKey 私钥解密文件
func DecryptTextWithRsaByPrivateKey(content []byte, privateKeyFileName string) (raw []byte, err error) {
	publicKeyRaw, err := ioutil.ReadFile(privateKeyFileName)
	if err != nil {
		return nil, err
	}
	block, _ := pem.Decode(publicKeyRaw)
	privateKey, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		return nil, err
	}
	return rsa.DecryptPKCS1v15(rand.Reader, privateKey, content)
}

func init() {
	rootCmd.AddCommand(rsaCmd)

	rsaCmd.Flags().StringVarP(&Text, "text", "t", "", "The text you want to encrypt or decrypt")
	rsaCmd.Flags().BoolVarP(&IsDecode, "decode", "d", false, "Indicates whether it is a decryption operation")
	rsaCmd.Flags().Int64VarP(&GenerateKeySize, "generate", "g", -1, "The rsa key size")
	rsaCmd.Flags().StringVarP(&PublicKey, "public", "u", "", "The public file name key of RSA")
	rsaCmd.Flags().StringVarP(&PrivateKey, "private", "r", "", "The private file name key of RSA")
}
